#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    : /counter_coruntine.py
@Time    : 2021-05-19 14:43:23
@Author  : Gray Snail
@Version : 1.0
@Contact : shtao2011@sina.com
@License : Copyright (c) 2021 Gray Snail
@Desc    : None
'''

# 解决 sleep 的阻塞问题

from collections import deque
import time
import heapq


class Scheduler(object):
    """
    通过两个队列实现了对 sleep 非阻塞和多任务调度

    思路：将长时间的阻塞转换为任务需要在某一时刻执行 time line
        task 按照 time line 排序
        在没有任务执行的时间段，调度器统一 sleep block


        yield from 从一个协程获取执行权


    asyncio 原理
        核心：类似的调度器

        ```python
        import asyncio

        async def main():
            print('Hello')
            await asyncio.sleep(0.2)
            print('World')

        asyncio.run(main())
        ```

        await 相当于 yield from

    """

    # 检测活跃任务并执行
    def __init__(self):
        self.ready = deque()
        self.current = None
        self.sequence = 0
        self.sleeping = []

    def sleep(self, delay):
        deadline = time.time() + delay
        # heap 排序
        self.sequence += 1
        heapq.heappush(self.sleeping, (deadline, self.sequence, self.current))

        self.current = None
        yield  # 交出控制权

    def new_task(self, coro):
        # coro 必须是生成器
        self.ready.append(coro)

    def run(self):
        while self.ready or self.sleeping:
            if not self.ready:
                deadline, _, coro = heapq.heappop(self.sleeping)
                delta = deadline - time.time()
                if delta > 0:
                    time.sleep(delta)
                self.ready.append(coro)

            self.current = self.ready.popleft()
            try:
                # send 属于生成器语法
                self.current.send(None)
                if self.current:
                    self.ready.append(self.current)

            except StopIteration:
                pass


scheduler = Scheduler()


def count_down(n):
    while n > 0:
        print('Down', n)
        yield from scheduler.sleep(1)
        n -= 1


def count_up(n):
    x = 0
    while x < n:
        print('Up', x + 1)
        yield from scheduler.sleep(0.5)
        x += 1


if __name__ == '__main__':

    s = time.time()
    scheduler.new_task(count_down(5))
    scheduler.new_task(count_up(15))

    scheduler.run()
    print(time.time() - s)
